from common_import import *


def date_to_num(date_str):
    return datetime.datetime.strptime(date_str, "%Y-%m-%d").toordinal()


def check_and_plot_sales(i):
    # 获取数据
    data = tool.get_np(f"product_daily_total/{i}.csv")
    if data is None:
        return
    # 将日期转换为数字
    dates = np.array([date_to_num(date) for date in data["date"]])
    quantities = np.array(data["total_quantity"], dtype=float)

    # 定义检查的日期范围
    start_date = date_to_num("2023-06-24")
    end_date = date_to_num("2023-06-30")

    # 找到2023-06-24到2023-06-30之间的日期索引
    mask = (dates >= start_date) & (dates <= end_date)
    # 计算该日期范围内的平均销售量
    if np.sum(mask) > 0:  # 确保日期范围内有数据
        avg_quantity = np.max(quantities[mask])
    else:
        avg_quantity = 0

    # 检查平均销售量是否大于2.5
    if avg_quantity > 2.5:
        # # 绘制三年以来的销售量曲线
        # plt.figure(figsize=(10, 6))
        # plt.plot(
        #     [datetime.datetime.fromordinal(int(d)) for d in dates],
        #     quantities,
        #     marker="o",
        # )
        # plt.title(f"Sales Quantity Over 3 Years for {mapping.get_name(i)}")
        # plt.xlabel("Date")
        # plt.ylabel("Total Quantity")
        # plt.grid(True)
        # mydraw.show_or_print(f"pick_picture/{i}.png")
        return 1
    return 0


def pick_product():
    result = []
    for i in range(1, 252):
        if check_and_plot_sales(i) == 1:
            result.append(i)
    print([mapping.get_name(i) for i in result])
    print(len(result))


if __name__ == "__main__":
    pick_product()
